package com.itheima.stargraph.service.impl;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.json.JSONUtil;
import com.itheima.stargraph.pojo.ComfyuiTask;
import com.itheima.stargraph.pojo.MessageBase;
import com.itheima.stargraph.service.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * Comfyui消息处理业务实现类
 *
 * @Author mr.wu
 * @Date 2025-6-12 11:40
 */
@Service
public class ComfyuiMessageServiceImpl implements ComfyuiMessageService {

    @Autowired
    private WsNotifyService wsNotifyService;

    /**
     * 类型：生图进度
     */
    private static final String TYPE_PROGRESS = "progress";

    /**
     * 类型：生图结果
     */
    private static final String TYPE_EXECUTED = "executed";

    /**
     * 类型：生图错误
     */
    private static final String TYPE_EXECUTION_ERROR = "execution_error";

    /**
     * 根据消息的类型不同，推送消息给客户端（前端）
     *
     * @param messageBase
     */
    @Override
    public void handleMsg(MessageBase messageBase) {
        String type = messageBase.getType();
        if(type.equals(TYPE_PROGRESS)) {
            handleProgress(messageBase);
        } else if(type.equals(TYPE_EXECUTED)) {
            handleExecuted(messageBase);
        } else if(type.equals(TYPE_EXECUTION_ERROR)) {
            handleExecutionError(messageBase);
        }
    }

    @Autowired
    private IUserFundService userFundService;


    /**
     * 生图错误消息推送给前端
     * @param messageBase
     */
    private void handleExecutionError(MessageBase messageBase) {
        HashMap<String, Object> data = messageBase.getData();
        data.put("type",TYPE_EXECUTION_ERROR);

        Object promptIdObj = data.getOrDefault("prompt_id", null);
        if(promptIdObj==null) {
            return;
        }
        String promptId = String.valueOf(promptIdObj);
        ComfyuiTask comfyuiTask = text2ImageService.getStartedTask(promptId);
        String clientId = comfyuiTask.getWsClientId();//前端的clientId

        //冻结积分归还
        userFundService.refund(comfyuiTask.getUserId(), Long.valueOf(comfyuiTask.getSize()));

        wsNotifyService.sendMsgToUser(clientId, JSONUtil.toJsonStr(data));
    }

    @Value("${comfyuiHost}")
    private String comfyuiHost;

    @Autowired
    private IUserResultService userResultService;

    /**
     * 生图结果消息推送给前端
     * @param messageBase
     */
    private void handleExecuted(MessageBase messageBase) {

        HashMap<String, Object> data = messageBase.getData();

        Map<String,Object> output = (Map<String, Object>) data.getOrDefault("output", null);
        if(CollUtil.isEmpty(output)){
            return;
        }

        List<Map<String,Object>> images = (List<Map<String, Object>>) output.getOrDefault("images",null);
        if(CollUtil.isEmpty(images)){
            return;
        }

        List<String> urls = images.stream().map(x -> {
            String urlTemplate = "http://%s/api/view?filename=%s&type=%s&subfolder=";
            String url = String.format(urlTemplate, comfyuiHost, x.get("filename"), x.get("type"));
            return url;
        }).collect(Collectors.toList());

        Map<String,Object> resultMap = new HashMap<>();
        resultMap.put("type","imageResult");
        resultMap.put("urls",urls);


        Object promptIdObj = data.getOrDefault("prompt_id", null);
        if(promptIdObj==null) {
            return;
        }
        String promptId = String.valueOf(promptIdObj);
        ComfyuiTask comfyuiTask = text2ImageService.getStartedTask(promptId);
        String clientId = comfyuiTask.getWsClientId();//前端的clientId

        //批量保存生图结果数据
        userResultService.saveBatch(comfyuiTask.getUserId(),urls);

        //积分扣除
        userFundService.deduct(comfyuiTask.getUserId(), Long.valueOf(comfyuiTask.getSize()));

        wsNotifyService.sendMsgToUser(clientId, JSONUtil.toJsonStr(resultMap) );
    }

    @Autowired
    private Text2ImageService text2ImageService;

    /**
     * 生图进度消息推送给前端
     * @param messageBase
     */
    private void handleProgress(MessageBase messageBase) {
        HashMap<String, Object> data = messageBase.getData();
        data.put("type",TYPE_PROGRESS);

        Object promptIdObj = data.getOrDefault("prompt_id", null);
        if(promptIdObj==null) {
            return;
        }
        String promptId = String.valueOf(promptIdObj);
        ComfyuiTask comfyuiTask = text2ImageService.getStartedTask(promptId);
        String clientId = comfyuiTask.getWsClientId();//前端的clientId
        wsNotifyService.sendMsgToUser(clientId, JSONUtil.toJsonStr(data));
    }
}
